(module fmap mzscheme
  (require (prefix list: (lib "list.ss"))
           (lib "unitsig.ss")
           (lib "etc.ss")
           (rename mzscheme mzscheme:map map)
           "ordered.ss"
           "set-signature.ss"
   )

  (provide fmap^
           instantiate-fmap
           use-fmap
           fmap@)

  (define-signature fmap^
    (empty
     insert
     insert-all
     assoc->fmap
     ->assoc
     get
     pairs
     remove
     remove-all
     member?
     map
     map-pairs
     map-keys
     map-values
     fold
     keys
     vals
     size
     fmap?))

  (define (instantiate-fmap set@ order@)
    (define FMAP@ (fmap@ set@))
    (compound-unit/sig 
      (import)
      (link (ALPHA      : ordered^ (order@))
            (ALPHA-HEAP : fmap^    (FMAP@ ALPHA)))
      (export (open ALPHA-HEAP))))
  
  
  (define-syntax (use-fmap stx)
    (syntax-case stx ()
      [(_ fmap-unit)
       (datum->syntax-object stx `(,#`define-values/invoke-unit/sig fmap^ ,#`fmap-unit))]
      [(_ fmap-unit prefix) 
       (datum->syntax-object stx `(,#`define-values/invoke-unit/sig fmap^ ,#`fmap-unit ,#`prefix))]))

  (define-struct entry (key value))
  
  (define (fmap@ set@)
    (unit/sig fmap^
      (import ordered^)
        
      (define (unwrap x) 
        (if (entry? x) (entry-key x) x))

      (define entry-order@
        (predicates->ordered (lambda (a b) (elm= (unwrap a) (unwrap b)))
                             (lambda (a b) (elm< (unwrap a) (unwrap b)))
                             (lambda (a b) (elm<= (unwrap a) (unwrap b)))))
        
      (use-set (instantiate-set set@ entry-order@) set)

      (define empty set:empty)

      (define (insert t k v) (set:insert t (make-entry k v)))

      (define (insert-all t lst)
        (list:foldl (lambda (p acc)
                      (insert acc (list:first p) (list:second p))) 
                    t lst))
        
      (define (assoc->fmap lst) (insert-all empty lst))

      (define ->assoc pairs)

      (define (fail-not-found fn-name k)
        (raise (make-exn:fail:contract
                (string->immutable-string (format "~a: no value found for key: ~s" fn-name k))
                (current-continuation-marks))))

      (define get
        (opt-lambda
            (t k [fail-thunk (lambda () (fail-not-found 'get k))])
          (if (not (set:member? t k))
              (fail-thunk)
              (entry-value (set:get t k)))))

      (define (pairs t)
        (mzscheme:map (lambda (e) (list (entry-key e) (entry-value e)))
                      (set:elements t)))

      (define remove set:remove)

      (define remove-all set:remove-all)

      (define member? set:member?)

      (define (map fn t)
        (set:fold
         (lambda (e acc) (cons (fn (entry-key e) (entry-value e)) acc))
         list:empty t))

      (define (map-pairs fn t)
        (set:fold
         (lambda (e acc)
           (let-values ([(new-k new-v) (fn (entry-key e) (entry-value e))])
             (set:insert acc new-k new-v)))
         set:empty t))

      (define (map-keys fn t)
        (set:fold
         (lambda (e acc) (set:insert acc (fn (entry-key e)) (entry-value e)))
         set:empty t))

      (define (map-values fn t)
        (set:fold
         (lambda (e acc) (set:insert acc (entry-key e) (fn (entry-value e))))
         set:empty t))

      (define (fold fn init t)
        (set:fold
         (lambda (e acc) (fn (entry-key e) (entry-value e) acc))
         init t))

      (define (keys t) (map (lambda (x y) x) t))

      (define (vals t) (map (lambda (x y) y) t))

      (define size set:size)

      (define fmap? set:set?))))
  
